from utils.generate_dataset import generate_dataset
from utils.log_output import simple_text_log, simple_metric_log
from PreRec.Recommender import Recommender
from TransFM.train import train
from TransFM.dataset import BookUserInfo
import torch
import pickle
import os
import json
from utils.upload_db import upload_db


def start_train():
    # Load hyper_params
    hyper_params_prerec = json.load(open(os.path.join('config', 'hyper_params_prerec.json'), 'rt'))
    hyper_params_transfm = json.load(open(os.path.join('config', 'hyper_params_transfm.json'), 'rt'))

    simple_text_log('train', 'Generating dataset...')

    # Step 1: Generate Dataset
    generate_dataset()

    simple_text_log('train', 'Pre-Recommendation...')
    # Step 2: PreRec
    prerec = Recommender(hyper_params_prerec)
    result = prerec.recommend()

    # Save result
    result_tensor = torch.empty([len(result), 500], dtype=torch.long)
    for i in range(len(result)):
        result_tensor[i] = torch.tensor(result[i])

    os.makedirs('prerec_result', exist_ok=True)
    torch.save(result_tensor, os.path.join('prerec_result/prerec_result.pkl'))

    simple_text_log('train', 'Loading book and user information...')
    # Step 3: Train full model
    book_user_info = BookUserInfo(hyper_params_transfm)
    hyper_params_transfm = book_user_info.update_hyper_params()
    simple_metric_log('train_metric', 'book_cnt', hyper_params_transfm['book_cnt'])
    simple_metric_log('train_metric', 'user_cnt', hyper_params_transfm['user_cnt'])

    simple_text_log('train', 'Training TransFM Model...')
    result = train(hyper_params_transfm, book_user_info)

    # Save the result
    os.makedirs('result', exist_ok=True)
    pickle.dump(result, open(os.path.join('result', 'result.pkl'), 'wb'))
    simple_text_log('train', 'Training finished.')

    # Upload db to frontend server
    upload_db()


if __name__ == '__main__':
    start_train()
